import pathlib
import shutil

import keras_autodoc

from autokeras import auto_model
from autokeras import task
from autokeras.hypermodel import block
from autokeras.hypermodel import head
from autokeras.hypermodel import hyperblock


# For each class to document, it is possible to:
# 1) Document only the class: [classA, classB, ...]
# 2) Document all its methods: [classA, (classB, '*')]
# 3) Choose which methods to document (methods listed as strings):
# [classA, (classB, ['method1', 'method2', ...]), ...]
# 4) Choose which methods to document (methods listed as qualified names):
# [classA, (classB, [module.classB.method1, module.classB.method2, ...]), ...]
PAGES = {
    'auto_model.md': [
        auto_model.AutoModel,
        auto_model.AutoModel.fit,
        auto_model.AutoModel.predict,
    ],
    'graph_auto_model.md': [
        'autokeras.GraphAutoModel',
        'autokeras.GraphAutoModel.fit',
        'autokeras.GraphAutoModel.predict',
    ],
    'block.md': [
         hyperblock.ImageBlock,
         hyperblock.TextBlock,
         hyperblock.StructuredDataBlock,
         block.ResNetBlock,
         block.XceptionBlock,
         block.ConvBlock,
         block.RNNBlock,
         block.Merge
     ],
    'task.md': [
        task.ImageClassifier,
        task.ImageRegressor,
        task.TextClassifier,
        task.TextRegressor
    ],
    'head.md': [head.ClassificationHead, head.RegressionHead]
}


ROOT = 'http://autokeras.com/'

autokeras_dir = pathlib.Path(__file__).resolve().parents[1]


def generate(dest_dir):
    template_dir = autokeras_dir / 'docs' / 'templates'
    doc_generator = keras_autodoc.DocumentationGenerator(
        PAGES,
        'https://github.com/keras-team/autokeras/blob/master',
        template_dir,
        autokeras_dir / 'examples',
    )
    doc_generator.generate(dest_dir)
    readme = (autokeras_dir / 'README.md').read_text()
    index = (template_dir / 'index.md').read_text()
    index = index.replace('{{autogenerated}}', readme[readme.find('##'):])
    (dest_dir / 'index.md').write_text(index, encoding='utf-8')
    shutil.copyfile(autokeras_dir / '.github' / 'CONTRIBUTING.md',
                    dest_dir / 'contributing.md')


if __name__ == '__main__':
    generate(autokeras_dir / 'docs' / 'sources')
